# import scheduler
import const
import functools
import threading
import planner

from channel import Message
from world import World, BlockStack
from block import Section
from planner import Pred

import random

class BlockGoal:
	"""The goal is to find a block that fits the 'slot', in order to
		trade it with the 'agent' and receive the block that will then be traded
		for the 'replacement'. This block will then be used to satisfy the
		goal that this goal 'depends' on."""
	def __init__(self, slot=None, name=None, block_id=None, depends=None):
		self.slot = slot
		self.agent = name
		self.replacement = block_id
		self.depends = depends

class Query:
	def __init__(self):
		self.condition = threading.Condition()

	def acquire(self):
		self.condition.acquire()

	def release(self):
		self.condition.release()
	
	def wait(self):
		self.condition.wait()
	
	def reply(self, answer, answer2 = None):
		self.condition.acquire()
		self.answer = answer
		self.answer2 = answer2
		self.condition.notifyAll()
		self.condition.release()

class Agent:
	"""An agent."""

	def __init__(self, name, world, sections):
		self.name = name
		self.world = world
		self.goalstacks = sections
		self.usedBlocks = {}
		self.pendingBlockQuery = None	# query to an owner of a block
		self.pendingSectionQuery = None	# query to an owner of a section
		self.commitments = []
		self.commitLock = threading.Lock()
		self.usedLock = threading.Lock()
		self.choosePhaseMutex = threading.Semaphore(0)
		self.initialSlots = 0 # number of slots before making commitments
		self.overCount = 0 # how many agents finished the choosing phase
		for goal in self.goalstacks:
			self.initialSlots += len(goal.slots)

	def setChannel(self, channel):
		"""Designates the communication channel that this agent will use for sending and receiving messages"""
		self.channel = channel

	def findBlock(self, slot):
		goal = BlockGoal(slot)
		blacklist = []
		block_id = self.findBlockAux(goal, blacklist, 0)
		return block_id

	def claim(self, block_id, old_owner=None):
		# try to set ourselves as the owner of this block
		if old_owner is not None:
			ret = self.world.changeOwner(block_id, old_owner, self.name)
		else:
			ret = self.world.setOwner(block_id, self.name)
		if not ret: 
			# apparently someone beat us to it
			self.printLocked(self.name + "'s claim on block " +
						str(block_id) + " failed")
			return ret
		# in case of success, let everyone know
		s = self.name + " claims " + str(block_id) 
		if old_owner is not None:
			s += " from " + old_owner
		self.printLocked(s)
		msg = Message(const.CHOOSE, block_id)
		self.channel.bcastMessage(self.name, msg)
		return True
	
	def findBlockAux(self, goal, blacklist, depth):

		# there are no more unowned blocks
		if not filter(lambda x: x.owner is None, self.world.getBlocks()):
			return None
		# we reached maximum depth
		if depth > const.MAX_DEPTH:
			return None

		# try to find block in pool 
		blocks = self.world.findMatchingBlocks(goal.slot)
		random.shuffle(blocks)
		if blocks:
			for b in blocks:
				self.usedLock.acquire()
				if self.claim(b.id):
					# we will release the lock in satisfyGoal
					return self.satisfyGoal(goal, b.id)
				self.usedLock.release()
		else:
		# try to find block at other agents
			blocks = self.world.findAllMatchingBlocks(goal.slot)
			random.shuffle(blocks)
			for b in blocks:
				# this block is already involved in trading
				if b.id in blacklist:
					continue
				# find and commit the slot that this block fills
				slot = self.findSlot(b)
				if not slot:
					continue
				# the new goal is to find a block that can fill this slot
				new_goal = BlockGoal(slot, b.owner, b.id, goal)  

				# call findBlockAux recursively to find the block
				blacklist.append(b.id)
				block = self.findBlockAux(new_goal, blacklist, depth+1)
				blacklist.pop()

				# if we found the block, use it to trade
				if block is not None:
					return self.satisfyGoal(goal, block)

				# if not, decommit the slot and return
				msg = Message(const.SLOT_DECOMMIT, b.id) 
				self.channel.sendMessage(self.name, b.owner, msg) 

	def findSlot(self, block):
		self.pendingBlockQuery = Query()

		# we talk to the owner of the block 'block'
		self.pendingBlockQuery.acquire()
		msg = Message(const.SLOT_ASK, block.id) 
		self.channel.sendMessage(self.name, block.owner, msg)
		self.pendingBlockQuery.wait()
		self.pendingBlockQuery.release()

		# they tell us what slot the block is for
		return self.pendingBlockQuery.answer

	def satisfyGoal(self, goal, block_id):
		# if this is the root goal
		if goal.depends == None:
			return block_id

		# this is an intermediate goal
		self.usedLock.release()

		# tell te agent to own the block we found
		if self.name != goal.agent:
			msg = Message(const.SLOT_TRADE, goal.replacement, block_id)
			self.channel.sendMessage(self.name, goal.agent, msg)
			# we own the replacement block: claim and mark as used atomically
			self.usedLock.acquire()
			self.claim(goal.replacement, goal.agent)
		else:
		# switch blocks
			self.printLocked(self.name + " is switching blocks "
				+ str(goal.replacement) + " and " + str(block_id))
			self.usedLock.acquire()
			self.usedBlocks[block_id] = self.usedBlocks[goal.replacement]
			del self.usedBlocks[goal.replacement]
			self.commitLock.acquire()
			self.cancelCommitment(const.KEEP_BLOCK, self.name, goal.replacement)
			self.commitLock.release()
		# return the block we found or got through trade
		return goal.replacement

	def shareBlock(self, our_section, merged_slot):
		# lock all of our blocks in that section
		self.usedLock.acquire()
		self.commitLock.acquire()
		mySlotDict = {}
		usedSlots = dict([(v, k) for (k, v) in self.usedBlocks.items()])
		for slot in our_section.slots:
			if slot in usedSlots:
				b = usedSlots[slot]
				self.makeCommitment(const.KEEP_BLOCK, self, b)
				mySlotDict[slot] = b
		self.usedLock.release()
		self.commitLock.release()

		# get all blocks matching our slot
		blocks = self.world.findAllMatchingBlocks(merged_slot)
		random.shuffle(blocks)

		# for each block, try to share the section containing it with its owner
		for block in blocks:
			if block.owner == self.name:
				continue
				
			self.pendingSectionQuery = Query()

			# get the section
			self.pendingSectionQuery.acquire()
			msg = Message(const.SECTION_ASK, block.id) 
			self.channel.sendMessage(self.name, block.owner, msg)
			self.pendingSectionQuery.wait()
			self.pendingSectionQuery.release()

			their_section = self.pendingSectionQuery.answer
			theirSlotDict = self.pendingSectionQuery.answer2

			if their_section is None:
				continue

			# initialize used block lists
			my_blocks = mySlotDict.values()
			their_blocks = theirSlotDict.values()
			new_blocks = []
			
			their_slot = None
			for slot in theirSlotDict:
				if theirSlotDict[slot] == block.id:
					their_slot = slot

			merged_section = Section.merge(our_section, their_section,
							merged_slot, their_slot)

			if merged_section is None:
				# decommit remotely
				msg = Message(const.SECTION_DECOMMIT, their_section)
				self.channel.sendMessage(self.name, block.owner, msg)
				filledAll = False
				continue

			mergedDict = {}

			filledAll = True
			for slot in merged_section.slots:
				found = False
				# look for a block in available blocks
				for b in my_blocks + their_blocks:
					b_block = self.world.getBlock(b)
					if slot.fits(b_block):
						if b in my_blocks:
							my_blocks.remove(b)
						if b in their_blocks:
							their_blocks.remove(b)
						mergedDict[slot] = b
						found = True
						# we found a block
						break
				if found:
					continue
				# look for a block anywhere in the world
				b = self.findBlock(slot)
				if b is not None:
					mergedDict[slot] = b
					new_blocks.append(b)
					self.usedBlocks[b] = slot 
					self.usedLock.release()
					continue
				# we failed

				# disown all blocks we owned for this
				for b in new_blocks:
					self.world.changeOwner(b, self, None)

				# mark blocks as unused
				for b in self.usedBlocks.keys():
					if self.usedBlocks[b] in merged_section.slots:
						del self.usedBlocks[b]
						self.printLocked("Agent " +  self.name + " deleting block " + str(b))
					
				# decommit remotely
				msg = Message(const.SECTION_DECOMMIT, their_section)
				self.channel.sendMessage(self.name, block.owner, msg)
				filledAll = False
				break

			if not filledAll:
				continue 

			# decommit the unused blocks
			self.printLocked(self.name + " deleting " + str(my_blocks))
			for b in my_blocks:
				self.commitLock.acquire()
				self.cancelCommitment(const.KEEP_BLOCK, self.name, b)
				self.commitLock.release()
				self.world.changeOwner(b, self, None)
				del self.usedBlocks[b]
				self.printLocked("Agent " +  self.name + " deleting block " + str(b))

			# mark used blocks
			self.usedLock.acquire()
			for slot in mergedDict:
				b = mergedDict[slot]
				if b in self.usedBlocks:
					del self.usedBlocks[b]
					self.printLocked("Agent " +  self.name + " deleting block " + str(b))
				self.usedBlocks[b] = slot
			self.usedLock.release()

			printstr = self.name + " and " + block.owner + \
			"\n\tmerged sections " + str(map(str, our_section.slots)) + \
			"\n\tand " + str(map(str, their_section.slots)) + \
			"\n\tinto section " + str(map(str, merged_section.slots)) + \
			"\n\tusing blocks "

			for (ss, bb) in mergedDict.items():
				printstr += "\n\t\t" + str(ss) + " " +\
					str(self.world.getBlock(bb).props) + " " + str(bb)
			self.printLocked(printstr)

			# compose the trade message
			msg = Message(const.SECTION_TRADE, their_section, merged_section,\
																mergedDict)
			self.channel.sendMessage(self.name, block.owner, msg)

			# commit shared section 
			self.commitLock.acquire()
			self.makeCommitment(const.SHARE_SECTION, our_section, block.owner)
			self.commitLock.release()

			our_section.slots = merged_section.slots
			our_section.shared = block.owner
			return True

		# we couldn't find a block: decommit all our blocks
		self.commitLock.acquire()
		for b in mySlotDict.values():
			self.cancelCommitment(const.KEEP_BLOCK, self.name, b)
		self.commitLock.release()
		return False
					
	def chooseBlocks(self):
		"""Starts the block choosing step."""
		for goal in self.goalstacks:
			for slot in goal.slots:
				if slot not in self.usedBlocks.values():
					block_id = self.findBlock(slot)
					# save the block-slot association
					if block_id is not None:
						self.usedBlocks[block_id] = slot 
						self.usedLock.release()
					else:
						# try to find a section to use in common
						if self.shareBlock(goal, slot):
							# section has changed
							goalDone = False
							break

	def checkCommitment(self, type, block_id):
		for c in self.commitments:
			if c.type == type and c.block_id == block_id: 
				return True
		return False

	def makeCommitment(self, type, agent, block_id):
		for c in self.commitments:
			if c.type == type and c.agent == agent and c.block_id == block_id: 
				return
		self.commitments.append(Commitment(type, agent, block_id))

	def cancelCommitment(self, type, agent, block_id):
		for c in self.commitments:
			if c.type == type and c.agent == agent and c.block_id == block_id:
				self.commitments.remove(c)
				return

	def printLocked(self, str):
		self.channel.printLock.acquire()
		print str
		self.channel.printLock.release()

	
	def recvMessage(self, sender, message):
		"""Function called upon receival of a message from another agent."""

		if message.type == const.SLOT_ASK:
			# someone wants to know what slot is a block that we own filling
			self.usedLock.acquire()
			count = self.usedBlocks.keys().count(message.arg1)
			if count == 0:
				# self.printLocked(self.name + " says " + str(message.arg1)
				#			+ " not in used blocks")
				answer = None
				self.usedLock.release()
			else:
				answer = self.usedBlocks[message.arg1]
				self.usedLock.release()
				# check if we have committed this block to another agent
				self.commitLock.acquire()
				if self.checkCommitment(const.KEEP_BLOCK, message.arg1):
					# self.printLocked(self.name 
					#	+ " says it has a previous commitment for " 
					#	+ str(message.arg1))
					answer = None
				else:
					self.makeCommitment(const.KEEP_BLOCK, sender, message.arg1)
				self.commitLock.release()
			msg = Message(const.SLOT_ANSWER, answer)
			self.channel.sendMessage(self.name, sender, msg)

		if message.type == const.SLOT_ANSWER:
			self.pendingBlockQuery.reply(message.arg1)
			
		if message.type == const.SLOT_TRADE:
			if sender != self.name:
				self.usedLock.acquire()
				self.claim(message.arg2, sender)
				self.usedBlocks[message.arg2] = self.usedBlocks[message.arg1] 
				del self.usedBlocks[message.arg1]
				self.usedLock.release()
			self.commitLock.acquire()
			self.cancelCommitment(const.KEEP_BLOCK, sender, message.arg1)
			self.commitLock.release()

		if message.type == const.SLOT_DECOMMIT:
			self.commitLock.acquire()
			self.cancelCommitment(const.KEEP_BLOCK, sender, message.arg1)
			self.commitLock.release()

		if message.type == const.SECTION_ASK:
			msg_none = Message(const.SECTION_ANSWER, None)

			self.usedLock.acquire();
			# find the slot that the block is occupying
			count = self.usedBlocks.keys().count(message.arg1)
			if count == 0:
				# self.printLocked(self.name + " says " + str(message.arg1)
				#			+ " not in used blocks")
				self.usedLock.release()
				self.channel.sendMessage(self.name, sender, msg_none)
				return

			# find the section where we're using the block
			slot = self.usedBlocks[message.arg1]
			section = None
			for goal in self.goalstacks:
				if slot in goal.slots:
					section = goal
			if section == None:
				# the slot isn't in any section - wtf?
				# self.printLocked(self.name + " says " + str(message.arg1)
				#		+ " not in any section")
				self.usedLock.release()
				self.channel.sendMessage(self.name, sender, msg_none)
				return

			# check if we have committed any block in this section 
			usedSlots = dict([(v, k) for (k, v) in self.usedBlocks.items()])
			for slot in section.slots:  
				if slot in usedSlots:
					block = usedSlots[slot]
					self.commitLock.acquire()
					commit = self.checkCommitment(const.KEEP_BLOCK, block)
					self.commitLock.release()
					if commit:
						# self.printLocked(self.name 
						#	+ " says it has a previous commitment for " 
						#	+ str(message.arg1))
						self.usedLock.release()
						self.channel.sendMessage(self.name, sender, msg_none)
						return

			# make commitment for all the blocks and the section itself
			slotDict = {}
			self.commitLock.acquire()
			for slot in section.slots:  
				if slot in usedSlots:
					block = usedSlots[slot]
					self.makeCommitment(const.KEEP_BLOCK, sender, block)
					slotDict[slot] = block
			self.makeCommitment(const.KEEP_SECTION, sender, section)
			self.commitLock.release()
	
			self.usedLock.release()

			msg = Message(const.SECTION_ANSWER, section, slotDict)
			self.channel.sendMessage(self.name, sender, msg)

			return

		if message.type == const.SECTION_ANSWER:
			self.pendingSectionQuery.reply(message.arg1, message.arg2)
			return

		if message.type == const.SECTION_TRADE:
			section = message.arg1
			new_section = message.arg2
			mergedDict = message.arg3

			usedSlots = dict([(v, k) for (k, v) in self.usedBlocks.items()])
			used_blocks = mergedDict.values()

			# decommit unused blocks and then the section
			self.usedLock.acquire()
			self.commitLock.acquire()
			for slot in section.slots:  
				if slot in usedSlots.keys():
					block = usedSlots[slot]
					del self.usedBlocks[block]
					self.printLocked("!!! Agent " +  self.name + " deleting block " + str(block))
					if block not in used_blocks:
						self.cancelCommitment(const.KEEP_BLOCK, sender, block)
						self.world.changeOwner(block, self.name, None)
			self.cancelCommitment(const.KEEP_SECTION, sender, section)
			self.commitLock.release()

			# update used blocks
			for slot in mergedDict:
				b = mergedDict[slot]
				if b in self.usedBlocks:
					del self.usedBlocks[b]
				self.usedBlocks[b] = slot
			self.usedLock.release()

			# commit shared section
			self.makeCommitment(const.SHARE_SECTION, sender, section)

			# change goal to reflect sharing
			section.slots = new_section.slots
			section.shared = sender

			return

		if message.type == const.SECTION_DECOMMIT:

			section = message.arg1
			usedSlots = dict([(v, k) for (k, v) in self.usedBlocks.items()])

			# decommit blocks
			self.usedLock.acquire()
			for slot in section.slots:  
				if slot in usedSlots:
					block = usedSlots[slot]
					self.cancelCommitment(const.KEEP_BLOCK, sender, block)
			self.usedLock.release()

			# decommit the section  
			self.cancelCommitment(const.KEEP_SECTION, sender, section)
			return

		if message.type == const.CHOOSE_PHASE_OVER:
			self.overCount += 1
			if self.overCount == self.channel.agentCount() - 1:
				self.choosePhaseMutex.release()

	def countInitialSlots(self):
		return self.initialSlots

	def countTotalSlots(self):
		slots = 0
		for goal in self.goalstacks:
			slots += len(goal.slots)
		return slots
	
	def countMatchedSlots(self):
		return len(self.usedBlocks)

	def countSharedSlots(self):
		sharedSlots = 0
		for goal in self.goalstacks:
			if goal.shared is not None:
				sharedSlots += len(goal.slots)
		return sharedSlots

	def makeGoals(self):
		# process the goal stacks and build the goal states
		goals = []
		usedSlots = dict([(v, k) for (k, v) in self.usedBlocks.items()])
		for goal in self.goalstacks:
			isInstantiated = True
			for slot in goal.slots:
				if slot not in usedSlots:
					isInstantiated = False
					break
			if not isInstantiated:
				break
			prev = usedSlots[goal.slots[0]]
			for slot in goal.slots[1:]:
				goals.append(Pred(const.ON, usedSlots[slot], prev))
				prev = usedSlots[slot]
		return goals
	
	def execute(self):
		"""Starts the planning/execution step."""
		goals = self.makeGoals()
		worldState = self.world.worldState()


		plan = planner.plan(goals, worldState)
		steps = plan.sort()

		#self.printLocked(self.name+" " + str(map(str,goals)))
		#self.channel.printLock.acquire()
		#if steps is not None:
		#	plan.print_tree()
		#	#print map(str,worldState)
		#else:
		#	print "no plan"
		#print "\n", self.name, str(map(str,goals))
		#self.channel.printLock.release()

		s = "\n" + self.name + "\n" + str(map(str,goals)) + "\n"
		if steps is not None:
			s += str(map(str,steps))
		else:
			s += "no plan"
		#s += "\n" + str(map(str, worldState))
		self.printLocked(s)
		return

	def runAgent(self):
		self.chooseBlocks()

		# end the block choosing phase and start the execution phase
		self.channel.bcastMessage(self.name, Message(const.CHOOSE_PHASE_OVER))
		self.choosePhaseMutex.acquire()
	
		self.execute()

class Commitment:
	"""A commitment made to or by another agent."""
	def __init__(self, type, agent, block_id):
		self.type = type
		self.agent = agent
		self.block_id = block_id
		return
